# -*- coding: utf-8 -*-
"""
基础算法抽象类
为所有彩票预测算法提供统一接口
"""

from abc import ABC, abstractmethod
import torch
import numpy as np
from typing import Dict, List, Tuple, Any, Optional
import pandas as pd

class BaseAlgorithm(ABC):
    """
    彩票预测算法基础抽象类
    所有具体算法都需要继承这个类并实现必要的方法
    """
    
    def __init__(self, algorithm_name: str, lottery_type: str, config: Dict[str, Any] = None):
        """
        初始化算法
        
        Args:
            algorithm_name: 算法名称
            lottery_type: 彩票类型 ('ssq' 或 'dlt')
            config: 算法配置参数字典
        """
        self.algorithm_name = algorithm_name
        self.lottery_type = lottery_type
        self.config = config or {}
        
        # 根据彩票类型设置基本参数
        if lottery_type == 'ssq':
            self.red_ball_range = (1, 33)
            self.blue_ball_range = (1, 16)
            self.red_ball_count = 6
            self.blue_ball_count = 1
        elif lottery_type == 'dlt':
            self.red_ball_range = (1, 35)
            self.blue_ball_range = (1, 12)
            self.red_ball_count = 5
            self.blue_ball_count = 2
        elif lottery_type == 'plw':
            # 排列5特殊处理：5个位置，每个位置0-9的数字
            self.red_ball_range = (0, 9)
            self.blue_ball_range = (0, 0)  # 排列5没有蓝球
            self.red_ball_count = 5
            self.blue_ball_count = 0
        elif lottery_type == 'kl8':
            # 快乐8特殊处理：20个位置，每个位置1-80的数字
            self.red_ball_range = (1, 80)
            self.blue_ball_range = (0, 0)  # 快乐8没有蓝球
            self.red_ball_count = 20
            self.blue_ball_count = 0
        elif lottery_type == 'fc3d':
            # 3D彩票特殊处理：3个位置，每个位置0-9的数字
            self.red_ball_range = (0, 9)
            self.blue_ball_range = (0, 0)  # 3D彩票没有蓝球
            self.red_ball_count = 3
            self.blue_ball_count = 0
        else:
            raise ValueError(f"不支持的彩票类型: {lottery_type}")
            
        self.is_trained = False
        self.model_weights = {}
        
    @abstractmethod
    def train(self, historical_data: pd.DataFrame, **kwargs) -> Dict[str, Any]:
        """
        训练算法模型
        
        Args:
            historical_data: 历史开奖数据
            **kwargs: 其他训练参数
            
        Returns:
            训练结果字典，包含损失、准确率等指标
        """
        pass
    
    @abstractmethod
    def predict(self, recent_data: np.ndarray, num_predictions: int = 1) -> List[Dict[str, List[int]]]:
        """
        生成预测结果
        
        Args:
            recent_data: 最近的开奖数据
            num_predictions: 预测组数
            
        Returns:
            预测结果列表，每个元素包含 {'red_balls': [], 'blue_balls': []}
        """
        pass
    
    @abstractmethod
    def calculate_pattern_weights(self, historical_data: pd.DataFrame) -> Dict[str, float]:
        """
        计算模式权重
        
        Args:
            historical_data: 历史开奖数据
            
        Returns:
            模式权重字典
        """
        pass
    
    def get_algorithm_info(self) -> Dict[str, Any]:
        """
        获取算法信息
        
        Returns:
            算法信息字典
        """
        return {
            'name': self.algorithm_name,
            'lottery_type': self.lottery_type,
            'is_trained': self.is_trained,
            'red_ball_range': self.red_ball_range,
            'blue_ball_range': self.blue_ball_range,
            'red_ball_count': self.red_ball_count,
            'blue_ball_count': self.blue_ball_count,
            'config': self.config
        }
    
    def validate_predictions(self, predictions: List[Dict[str, List[int]]]) -> List[Dict[str, List[int]]]:
        """
        验证并修正预测结果
        
        Args:
            predictions: 原始预测结果
            
        Returns:
            验证后的预测结果
        """
        validated_predictions = []
        
        for pred in predictions:
            red_balls = pred.get('red_balls', [])
            blue_balls = pred.get('blue_balls', [])
            
            # 验证红球
            red_balls = self._validate_red_balls(red_balls)
            
            # 验证蓝球
            blue_balls = self._validate_blue_balls(blue_balls)
            
            validated_predictions.append({
                'red_balls': red_balls,
                'blue_balls': blue_balls
            })
            
        return validated_predictions
    
    def _validate_red_balls(self, red_balls: List[int]) -> List[int]:
        """
        验证红球预测结果
        
        Args:
            red_balls: 红球预测结果
            
        Returns:
            验证后的红球结果
        """
        # 确保红球在有效范围内
        red_balls = [max(self.red_ball_range[0], min(self.red_ball_range[1], ball)) 
                    for ball in red_balls]
        
        # 去重并补足数量
        red_balls = list(set(red_balls))
        while len(red_balls) < self.red_ball_count:
            candidate = np.random.randint(self.red_ball_range[0], self.red_ball_range[1] + 1)
            if candidate not in red_balls:
                red_balls.append(candidate)
        
        # 保持指定数量并排序
        return sorted(red_balls[:self.red_ball_count])
    
    def _validate_blue_balls(self, blue_balls: List[int]) -> List[int]:
        """
        验证蓝球预测结果
        
        Args:
            blue_balls: 蓝球预测结果
            
        Returns:
            验证后的蓝球结果
        """
        # 确保蓝球在有效范围内
        blue_balls = [max(self.blue_ball_range[0], min(self.blue_ball_range[1], ball)) 
                     for ball in blue_balls]
        
        # 去重并补足数量
        blue_balls = list(set(blue_balls))
        while len(blue_balls) < self.blue_ball_count:
            candidate = np.random.randint(self.blue_ball_range[0], self.blue_ball_range[1] + 1)
            if candidate not in blue_balls:
                blue_balls.append(candidate)
        
        # 保持指定数量并排序
        return sorted(blue_balls[:self.blue_ball_count])
    
    def save_model(self, filepath: str) -> bool:
        """
        保存模型权重
        
        Args:
            filepath: 保存路径
            
        Returns:
            是否保存成功
        """
        try:
            torch.save({
                'algorithm_name': self.algorithm_name,
                'lottery_type': self.lottery_type,
                'config': self.config,
                'model_weights': self.model_weights,
                'is_trained': self.is_trained
            }, filepath)
            return True
        except Exception as e:
            print(f"保存模型失败: {e}")
            return False
    
    def load_model(self, filepath: str) -> bool:
        """
        加载模型权重
        
        Args:
            filepath: 模型文件路径
            
        Returns:
            是否加载成功
        """
        try:
            checkpoint = torch.load(filepath, map_location='cpu')
            self.algorithm_name = checkpoint['algorithm_name']
            self.lottery_type = checkpoint['lottery_type']
            self.config = checkpoint['config']
            self.model_weights = checkpoint['model_weights']
            self.is_trained = checkpoint['is_trained']
            return True
        except Exception as e:
            print(f"加载模型失败: {e}")
            return False


class PatternAnalyzer:
    """
    模式分析器基础类
    用于分析彩票数据中的各种模式
    """
    
    def __init__(self, lottery_type: str):
        self.lottery_type = lottery_type
    
    def analyze_consecutive_patterns(self, historical_data: pd.DataFrame, 
                                   pattern_length: int = 3) -> Dict[str, Any]:
        """
        分析连续号码模式
        
        Args:
            historical_data: 历史数据
            pattern_length: 模式长度
            
        Returns:
            连续模式分析结果
        """
        consecutive_patterns = {}
        
        # 分析红球连续模式
        red_columns = [col for col in historical_data.columns if 'Red_' in col or '红球_' in col]
        
        for period_len in range(2, pattern_length + 1):
            for start_idx in range(len(historical_data) - period_len + 1):
                period_data = historical_data.iloc[start_idx:start_idx + period_len]
                
                # 提取所有红球号码
                all_numbers = []
                for _, row in period_data.iterrows():
                    numbers = [row[col] for col in red_columns if pd.notna(row[col])]
                    all_numbers.extend(numbers)
                
                # 计算连续出现的号码
                consecutive_count = self._count_consecutive_numbers(all_numbers)
                for num, count in consecutive_count.items():
                    if count >= period_len:
                        pattern_key = f"consecutive_{num}_{period_len}"
                        if pattern_key not in consecutive_patterns:
                            consecutive_patterns[pattern_key] = []
                        consecutive_patterns[pattern_key].append({
                            'start_period': start_idx,
                            'length': period_len,
                            'count': count
                        })
        
        return consecutive_patterns
    
    def analyze_diagonal_patterns(self, historical_data: pd.DataFrame) -> Dict[str, Any]:
        """
        分析斜连号模式
        
        Args:
            historical_data: 历史数据
            
        Returns:
            斜连号模式分析结果
        """
        diagonal_patterns = {
            'right_diagonal': [],  # 右斜：递增序列
            'left_diagonal': []    # 左斜：递减序列
        }
        
        red_columns = [col for col in historical_data.columns if 'Red_' in col or '红球_' in col]
        
        for i in range(len(historical_data) - 2):
            # 取连续三期数据
            three_periods = historical_data.iloc[i:i+3]
            
            # 检查右斜模式 (递增)
            right_diagonals = self._find_right_diagonal_sequences(three_periods, red_columns)
            diagonal_patterns['right_diagonal'].extend(right_diagonals)
            
            # 检查左斜模式 (递减)  
            left_diagonals = self._find_left_diagonal_sequences(three_periods, red_columns)
            diagonal_patterns['left_diagonal'].extend(left_diagonals)
        
        return diagonal_patterns
    
    def _count_consecutive_numbers(self, numbers: List[int]) -> Dict[int, int]:
        """计算连续出现的号码次数"""
        from collections import Counter
        return Counter(numbers)
    
    def _find_right_diagonal_sequences(self, three_periods: pd.DataFrame, 
                                     red_columns: List[str]) -> List[Dict]:
        """查找右斜序列 (递增)"""
        sequences = []
        
        for period_idx in range(3):
            if period_idx >= len(three_periods):
                continue
                
            period_numbers = [three_periods.iloc[period_idx][col] 
                            for col in red_columns 
                            if pd.notna(three_periods.iloc[period_idx][col])]
            
            for num in period_numbers:
                if (period_idx < 2 and 
                    self._check_number_in_period(three_periods.iloc[period_idx + 1], red_columns, num + 1) and
                    self._check_number_in_period(three_periods.iloc[period_idx + 2], red_columns, num + 2)):
                    
                    sequences.append({
                        'sequence': [num, num + 1, num + 2],
                        'type': 'right_diagonal',
                        'start_period': period_idx
                    })
        
        return sequences
    
    def _find_left_diagonal_sequences(self, three_periods: pd.DataFrame, 
                                    red_columns: List[str]) -> List[Dict]:
        """查找左斜序列 (递减)"""
        sequences = []
        
        for period_idx in range(3):
            if period_idx >= len(three_periods):
                continue
                
            period_numbers = [three_periods.iloc[period_idx][col] 
                            for col in red_columns 
                            if pd.notna(three_periods.iloc[period_idx][col])]
            
            for num in period_numbers:
                if (period_idx < 2 and num >= 3 and
                    self._check_number_in_period(three_periods.iloc[period_idx + 1], red_columns, num - 1) and
                    self._check_number_in_period(three_periods.iloc[period_idx + 2], red_columns, num - 2)):
                    
                    sequences.append({
                        'sequence': [num, num - 1, num - 2],
                        'type': 'left_diagonal',
                        'start_period': period_idx
                    })
        
        return sequences
    
    def _check_number_in_period(self, period_row: pd.Series, red_columns: List[str], 
                              target_number: int) -> bool:
        """检查指定号码是否在某期中出现"""
        period_numbers = [period_row[col] for col in red_columns if pd.notna(period_row[col])]
        return target_number in period_numbers